package com.voiceqsologger.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.voiceqsologger.config.AsrModeProperties;
import com.voiceqsologger.service.StreamingAsrService;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.nio.ByteBuffer;

/**
 * 面试辅导实时 WebSocket：
 * - 接收麦克风 PCM 流 → 流式 ASR
 * - 每个最终分段，调用 LLM 生成“回答提示”
 * - 追求低延迟：ASR 启动即推送 asr_progress，LLM 前后推送耗时
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class InterviewCoachHandler extends BinaryWebSocketHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final StreamingAsrService streamingAsrService;
    private final AsrModeProperties modeProperties;
    private final com.voiceqsologger.service.LlmInterviewCoachClient coachClient;

    @SneakyThrows
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        JsonNode node = objectMapper.readTree(message.getPayload());
        String type = node.path("type").asText();
        if ("start".equalsIgnoreCase(type)) {
            int sampleRate = node.path("sampleRate").asInt(16000);
            String lang = node.path("lang").asText("ZH");
            String mode = node.has("mode") ? node.path("mode").asText(null) : null;
            streamingAsrService.start(session.getId(), sampleRate, lang, mode);
            session.sendMessage(new TextMessage("{\"event\":\"started\"}"));
        } else if ("stop".equalsIgnoreCase(type)) {
            streamingAsrService.stop(session.getId());
            session.sendMessage(new TextMessage("{\"event\":\"stopped\"}"));
        }
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        ByteBuffer buffer = message.getPayload();
        byte[] chunk = new byte[buffer.remaining()];
        buffer.get(chunk);

        final long[] asrStart = new long[1];
        streamingAsrService.append(session.getId(), chunk,
                text -> {
                    try {
                        session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                                .put("event", "partial")
                                .put("text", text)
                                .toString()));
                    } catch (Exception e) { log.error("send partial error", e); }
                },
                finalText -> {
                    try {
                        // 推送最终文本
                        session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                                .put("event", "final")
                                .put("text", finalText)
                                .put("asrCostMs", (asrStart[0] > 0 ? (System.currentTimeMillis()-asrStart[0]) : 0))
                                .toString()));

                        // 异步调用 LLM 生成提示，尽可能快速
                        new Thread(() -> {
                            long t0 = System.currentTimeMillis();
                            String tips;
                            try {
                                tips = coachClient.suggest(finalText);
                            } catch (Exception e) {
                                tips = "获取提示失败: " + e.getMessage();
                            }
                            long cost = System.currentTimeMillis() - t0;
                            try {
                                session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                                        .put("event", "coach")
                                        .put("mode", modeProperties.getMode())
                                        .put("tips", tips == null ? "" : tips)
                                        .put("llmCostMs", cost)
                                        .toString()));
                            } catch (Exception ignore) {}
                        }).start();
                    } catch (Exception e) { log.error("send final error", e); }
                },
                () -> {
                    asrStart[0] = System.currentTimeMillis();
                    try {
                        session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                                .put("event", "asr_progress")
                                .put("status", "start")
                                .toString()));
                    } catch (Exception ignore) {}
                });
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        streamingAsrService.stop(session.getId());
        super.afterConnectionClosed(session, status);
    }
}


